#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 1;
class KernelScatterMax
{
public:
  __aicore__ inline KernelScatterMax() {}
  __aicore__ inline void Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, uint32_t blockLength,
                              uint32_t tileNum, uint32_t tileLength,
                              uint32_t lasttileLength,
                              uint32_t typeKey, uint32_t totalLength, uint32_t unitLength, uint32_t indicesCount)
  {
    ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
    this->typeKey = typeKey;
    this->totalLength = totalLength;
    
    this->blockLength = blockLength;
    this->tileNum =
        tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
    this->tileLength = tileLength / BUFFER_NUM;
    this->lasttileLength = lasttileLength;
    uint32_t ALIGN_NUM = 32;
    if(typeKey == 0) ALIGN_NUM = 8;
    if(typeKey == 1) ALIGN_NUM = 16;
    if(typeKey == 3) ALIGN_NUM = 8;
    if(lasttileLength % ALIGN_NUM == 0){
      this->lasttileLengthAlign = lasttileLength;
    }else{
      this->lasttileLengthAlign = (lasttileLength + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;
    }
    this->unitLength = unitLength;
    this->indicesCount = indicesCount;
    this->indices_addr = indices;
    varGm.SetGlobalBuffer((__gm__ DTYPE_VAR *)var,
                        this->blockLength);
    updatesGm.SetGlobalBuffer((__gm__ DTYPE_VAR *)updates,
                        this->blockLength);
    pipe.InitBuffer(inQueueV, BUFFER_NUM, this->tileLength * sizeof(DTYPE_VAR));
    pipe.InitBuffer(inQueueU, BUFFER_NUM, this->tileLength * sizeof(DTYPE_VAR));
    pipe.InitBuffer(outQueueV, BUFFER_NUM, this->tileLength * sizeof(DTYPE_VAR));
    if(typeKey == 2){
      pipe.InitBuffer(calcBuf, this->tileLength * sizeof(half));
      pipe.InitBuffer(calcBuf2, this->tileLength * sizeof(half));
    }
    
  }
  __aicore__ inline void Process()
  {
    int32_t loopCount = this->tileNum * BUFFER_NUM;
    for(int32_t i = 0; i < indicesCount; i++)
    for(int32_t j = 0; j < loopCount; j++)
    {
      CopyIn(i, j);
      Compute(i, j);
      CopyOut(i, j);
      // 使用空转，隔离前后两组
      LocalTensor<DTYPE_VAR> updatesLocal = inQueueU.AllocTensor<DTYPE_VAR>();
      inQueueU.EnQue(updatesLocal);
      updatesLocal = inQueueU.DeQue<DTYPE_VAR>();
      LocalTensor<DTYPE_VAR> yLocal = outQueueV.AllocTensor<DTYPE_VAR>();
      outQueueV.EnQue(yLocal);
      inQueueU.FreeTensor(updatesLocal);
      yLocal = outQueueV.DeQue<DTYPE_VAR>();
      outQueueV.FreeTensor(yLocal);
    }
  }

private:
  __aicore__ inline void CopyIn(int32_t id, int32_t progress)
  {
    LocalTensor<DTYPE_VAR> varLocal = inQueueV.AllocTensor<DTYPE_VAR>();
    LocalTensor<DTYPE_VAR> updatesLocal = inQueueU.AllocTensor<DTYPE_VAR>();

    int32_t varSatrtIndex = ((__gm__ int32_t *)indices_addr)[id] * unitLength;
    int32_t updatesSatrtIndex = id * unitLength;
    if(progress != this->tileNum - 1){
      DataCopy(varLocal[0], varGm[varSatrtIndex + progress * this->tileLength],
              this->tileLength);
      DataCopy(updatesLocal[0], updatesGm[updatesSatrtIndex + progress * this->tileLength],
              this->tileLength);
    }else{
      DataCopy(varLocal[0], varGm[varSatrtIndex + progress * this->tileLength],
              this->lasttileLengthAlign);
      DataCopy(updatesLocal[0], updatesGm[updatesSatrtIndex + progress * this->tileLength],
              this->lasttileLengthAlign);
    }
    inQueueV.EnQue(varLocal);
    inQueueU.EnQue(updatesLocal);
  }
  __aicore__ inline void Compute(int32_t id, int32_t progress)
  {
    if(typeKey == 2){
      LocalTensor<int8_t> varLocal = inQueueV.DeQue<int8_t>();
      LocalTensor<int8_t> updatesLocal = inQueueU.DeQue<int8_t>();
      LocalTensor<int8_t> yLocal = outQueueV.AllocTensor<int8_t>();
      LocalTensor<half> varLocal_h = calcBuf.Get<half>();
      LocalTensor<half> updatesLocal_h = calcBuf2.Get<half>();
      if (progress == this->tileNum - 1) {
        Cast(varLocal_h, varLocal, RoundMode::CAST_NONE, this->lasttileLengthAlign);
        Cast(updatesLocal_h, updatesLocal, RoundMode::CAST_NONE, this->lasttileLengthAlign);
        Max(varLocal_h, varLocal_h, updatesLocal_h, this->lasttileLengthAlign);
        Cast(yLocal, varLocal_h, RoundMode::CAST_NONE, this->lasttileLengthAlign);
        for(int i = lasttileLength; i<lasttileLengthAlign; i++){
          yLocal.SetValue(i, varLocal.GetValue(i));
        }
      }else{
        Cast(varLocal_h, varLocal, RoundMode::CAST_NONE, this->tileLength);
        Cast(updatesLocal_h, updatesLocal, RoundMode::CAST_NONE, this->tileLength);
        Max(varLocal_h, varLocal_h, updatesLocal_h, this->tileLength);
        Cast(yLocal, varLocal_h, RoundMode::CAST_NONE, this->tileLength);
      }
      outQueueV.EnQue(yLocal);
      inQueueV.FreeTensor(varLocal);
      inQueueU.FreeTensor(updatesLocal);
    }else{
      LocalTensor<DTYPE_VAR> varLocal = inQueueV.DeQue<DTYPE_VAR>();
      LocalTensor<DTYPE_VAR> updatesLocal = inQueueU.DeQue<DTYPE_VAR>();
      LocalTensor<DTYPE_VAR> yLocal = outQueueV.AllocTensor<DTYPE_VAR>();

      if (progress == this->tileNum - 1) {
        Max(yLocal, varLocal, updatesLocal, this->lasttileLengthAlign);
        for(int i = lasttileLength; i<lasttileLengthAlign; i++){
          yLocal.SetValue(i, varLocal.GetValue(i));
        }
      }else{
        Max(yLocal, varLocal, updatesLocal, this->tileLength);
      }
      outQueueV.EnQue(yLocal);
      inQueueV.FreeTensor(varLocal);
      inQueueU.FreeTensor(updatesLocal);
    }
  }
  __aicore__ inline void CopyOut(int32_t id, int32_t progress)
  {
    LocalTensor<DTYPE_VAR> yLocal = outQueueV.DeQue<DTYPE_VAR>();
    int32_t varSatrtIndex = ((__gm__ int32_t *)indices_addr)[id] * unitLength;
    if(progress != this->tileNum - 1)
      DataCopy(varGm[varSatrtIndex + progress * this->tileLength],
                yLocal[0], 
                this->tileLength);
    else
      DataCopy(varGm[varSatrtIndex + progress * this->tileLength],
                yLocal[0], 
                this->lasttileLengthAlign);
    outQueueV.FreeTensor(yLocal);
  }

private:
  TPipe pipe;
  TQue<QuePosition::VECIN, BUFFER_NUM> inQueueV, inQueueU;
  TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueV;
  TBuf<QuePosition::VECCALC> calcBuf, calcBuf2;
  GlobalTensor<DTYPE_VAR> varGm, updatesGm;
  GM_ADDR indices_addr;
  uint32_t blockLength;
  uint32_t tileNum;
  uint32_t tileLength;
  uint32_t lasttileLength;
  uint32_t lasttileLengthAlign;
  uint32_t typeKey;
  uint32_t totalLength;
  uint32_t unitLength;
  uint32_t indicesCount;
  // uint32_t masklen;
};

extern "C" __global__ __aicore__ void scatter_max(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR workspace, GM_ADDR tiling) {
  GET_TILING_DATA(tiling_data, tiling);
  // TODO: user kernel impl
  KernelScatterMax op;
  
  op.Init(var, indices, updates, tiling_data.blockLength,
          tiling_data.tileNum, tiling_data.tileLength,
          tiling_data.lasttileLength, tiling_data.typeKey,
          tiling_data.totalLength,
          tiling_data.unitLength, tiling_data.indicesCount);
  op.Process();
  // TODO: user kernel impl
}
#ifndef __CCE_KT_TEST__
void scatter_max_do(uint32_t blockDim, void* l2ctrl, void* stream,
                       uint8_t* var, uint8_t* indices, uint8_t* updates,
                       uint8_t* workspace, uint8_t* tiling) {
  scatter_max<<<blockDim, l2ctrl, stream>>>(var, indices, updates, workspace, tiling);
}
#endif
